// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (C) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: cred specific symbols
 * Author: Huawei OS Kernel Lab
 * Create: Thu Sep 20 16:34:03 2022
 */

#include <asm/uaccess.h>
#include <linux/cred.h>
#include <linux/slab.h>
#include <linux/mutex.h>
#include <linux/string.h>
#include <linux/init_task.h>
#include <linux/fs_struct.h>
#include <linux/mmu_notifier.h>
#include <asm/mmu_context.h>
#include <linux/hugetlb.h>
#ifdef CONFIG_KEYS
#include <linux/keyctl.h>
#endif

#include <lnxbase/lnxbase.h>

#ifdef ENABLE_CURRENT_MM
static void __mm_init(struct mm_struct *mm, struct task_struct *p)
{
	mm->mmap = NULL;
	mm->mm_rb = RB_ROOT;
	mm->vmacache_seqnum = 0;
#ifdef CONFIG_RSS_THRESHOLD
	mm->rss_threshold = 0;
#endif
	atomic_set(&mm->mm_users, 1);
	atomic_set(&mm->mm_count, 1);
	seqcount_init(&mm->write_protect_seq);
	mmap_init_lock(mm);
	INIT_LIST_HEAD(&mm->mmlist);
	mm->core_state = NULL;
	mm_pgtables_bytes_init(mm);
	mm->map_count = 0;
	mm->locked_vm = 0;
	atomic_set(&mm->has_pinned, 0);
	atomic64_set(&mm->pinned_vm, 0);
	memset(&mm->rss_stat, 0, sizeof(mm->rss_stat));
	spin_lock_init(&mm->page_table_lock);
	spin_lock_init(&mm->arg_lock);
	mm_init_cpumask(mm);
	RCU_INIT_POINTER(mm->exe_file, NULL);
	mmu_notifier_subscriptions_init(mm);
	init_tlb_flush_pending(mm);
#if defined(CONFIG_TRANSPARENT_HUGEPAGE) && !USE_SPLIT_PMD_PTLOCKS
	mm->pmd_huge_pte = NULL;
#endif
	hugetlb_count_init(mm);
#ifdef CONFIG_TASK_PROTECT_LRU
	mm->protect = 0;
#endif
	mm->flags = 0;
	mm->def_flags = 0;
	init_new_context(p, mm);
	mm->user_ns = &init_user_ns;
}

struct rb_mm_node {
	struct rb_node rb_node;
	unsigned int cnode_idx;
	struct mm_struct mm;
};

static struct rb_root mm_rbroot = RB_ROOT;
static struct liblinux_pal_mutex_t mm_rb_lock = LIBLINUX_MUTEX_INITIALIZER;
static struct rb_mm_node *rb_simple_get_init(struct rb_root *root, u64 cnode_idx,
					     struct task_struct *tsk)
{
	struct rb_node **p = &root->rb_node;
	struct rb_node *parent = NULL;
	struct rb_mm_node *entry = NULL;
	struct rb_mm_node *mm_node = NULL;
	unsigned long mm_size = 0;

	while (*p) {
		parent = *p;
		entry = rb_entry(parent, struct rb_mm_node, rb_node);

		if (cnode_idx < entry->cnode_idx)
			p = &(*p)->rb_left;
		else if (cnode_idx > entry->cnode_idx)
			p = &(*p)->rb_right;
		else {
			tsk->mm = &entry->mm;
			return entry;
		}
	}

	/* not found */
	mm_size = sizeof(struct rb_mm_node) + cpumask_size();
	mm_node = vmalloc(mm_size);
	if (mm_node == NULL) {
		pr_err("current mm init failed !\n");
		return NULL;
	}
	mm_node->cnode_idx = cnode_idx;
	__mm_init(&mm_node->mm, tsk);
	tsk->mm = &mm_node->mm;

	rb_link_node(&mm_node->rb_node, parent, p);
	rb_insert_color(&mm_node->rb_node, root);
	return mm_node;
}

static void __init_current_mm(struct task_struct *p, unsigned int cnode_idx)
{
	liblinux_pal_mutex_lock(&mm_rb_lock);
	rb_simple_get_init(&mm_rbroot, cnode_idx, p);
	liblinux_pal_mutex_unlock(&mm_rb_lock);
}
#endif

struct liblinux_cred {
	struct cred cred;
	unsigned long cnode_idx;
};

static unsigned long max_cspace_slots;
static struct liblinux_cred *cnode_idx_map;

#ifdef CONFIG_KEYS
static int __init cred_session_keyring_init(void)
{
	int i;
	struct key *liblinux_session_keyring = NULL;

	liblinux_session_keyring = (struct key *)keyring_alloc(
					"uid.liblinux", GLOBAL_ROOT_UID, GLOBAL_ROOT_GID,
					&init_cred, (KEY_POS_ALL & ~KEY_POS_SETATTR) | KEY_USR_ALL,
					KEY_ALLOC_UID_KEYRING | KEY_ALLOC_IN_QUOTA, NULL, NULL);
	if (IS_ERR(liblinux_session_keyring)) {
		pr_err("alloc keyring failed: %ld\n", PTR_ERR(liblinux_session_keyring));
		return PTR_ERR(liblinux_session_keyring);
	}
	init_cred.session_keyring = liblinux_session_keyring;
	init_cred.jit_keyring = KEY_REQKEY_DEFL_SESSION_KEYRING;

	(void)get_cred(&init_cred);

	for (i = 0; i < max_cspace_slots; i++) {
		cnode_idx_map[i].cred.session_keyring = liblinux_session_keyring;
		cnode_idx_map[i].cred.jit_keyring = KEY_REQKEY_DEFL_SESSION_KEYRING;
	}

	return 0;
}
core_initcall(cred_session_keyring_init);
#endif

void __init liblinux_cred_init(void)
{
	int i, ret;
	struct liblinux_cred ldk_cred = {
		.cred = init_cred,
		.cnode_idx = 0,
	};

	ret = lnxbase_max_cspace_slots_get(&max_cspace_slots);
	if (ret < 0)
		panic("get cspace slots config failed, ret=%d\n", ret);

	/*cnode_idx num may exceed max_cspace_slots, so cnode_idx_map should bigger than max_space*/
	max_cspace_slots = max_cspace_slots * 2;

	/*
	 * cnode_idx num may exceed max_cspace_slots, so cnode_idx_map should bigger than max_space.
	 * reserve 100 slots for the expanding cspaces.
	 */
	max_cspace_slots = max_cspace_slots + 100;

	/* life cyle same as devhost */
	cnode_idx_map = kcalloc(max_cspace_slots, sizeof(struct liblinux_cred), GFP_KERNEL);

	if (cnode_idx_map == NULL)
		panic("cnode_idx_map kcalloc failed\n");

	for (i = 0; i < max_cspace_slots; i++)
		memcpy(&cnode_idx_map[i], &ldk_cred, sizeof(ldk_cred));
}

struct cred *acquire_cred_by_cnode_idx(unsigned int cnode_idx)
{
	int i;
	struct cred *cred = NULL;
	struct liblinux_cred *ldk_cred;
	static struct liblinux_pal_mutex_t cred_lock = LIBLINUX_MUTEX_INITIALIZER;

	liblinux_pal_mutex_lock(&cred_lock);
	for (i = 0; i < max_cspace_slots; i++) {
		int pos = (cnode_idx + i) % max_cspace_slots;

		ldk_cred = &cnode_idx_map[pos];
		if (ldk_cred->cnode_idx == 0) {
			ldk_cred->cnode_idx = cnode_idx;
			cred = &ldk_cred->cred;
			break;
		} else if (ldk_cred->cnode_idx == cnode_idx) {
			cred = &ldk_cred->cred;
			break;
		}
	}
	liblinux_pal_mutex_unlock(&cred_lock);

	return cred;
}

void liblinux_enter(int pid, ...)
{
	va_list args;
	unsigned int cnode_idx;
	unsigned int uid;
	unsigned int gid;
	unsigned int suid;
	unsigned int sgid;
	unsigned int euid;
	unsigned int egid;
	unsigned int fsuid;
	unsigned int fsgid;
	int umask;
	struct cred *cred = NULL;

	set_fs(USER_DS);
	va_start(args, pid);
	uid = va_arg(args, unsigned int);
	gid = va_arg(args, unsigned int);
	suid = va_arg(args, unsigned int);
	sgid = va_arg(args, unsigned int);
	euid = va_arg(args, unsigned int);
	egid = va_arg(args, unsigned int);
	fsuid = va_arg(args, unsigned int);
	fsgid = va_arg(args, unsigned int);
	umask = va_arg(args, int);
	cnode_idx = va_arg(args, unsigned int);

	va_end(args);
	current->pid = pid;
	current->tgid = pid;

	cred = acquire_cred_by_cnode_idx(cnode_idx);
	BUG_ON(cred == NULL);

	/* Maybe suid is not match, when thread context changed. but never hanppend now.
	 * It will fixed after linux cred imported.
	 */
	cred->uid.val = uid;
	cred->gid.val = gid;
	cred->suid.val = suid;
	cred->sgid.val = sgid;
	cred->euid.val = euid;
	cred->egid.val = egid;
	cred->fsuid.val = fsuid;
	cred->fsgid.val = fsgid;

	current->cred = cred;
	current->fs->umask = umask;
	clear_tsk_thread_flag(current, TIF_SIGPENDING);
	__set_current_state(TASK_RUNNING);
#ifdef ENABLE_CURRENT_MM
	__init_current_mm(current, cnode_idx);
#endif
}
